#******************************************************************************
#******************************************************************************

# import libraries

import math

from scipy.optimize import fsolve
        
from ..pipes import fic # ugly import but it works

from ..common import formulas as core

from dataclasses import dataclass

#******************************************************************************
#******************************************************************************
#******************************************************************************
#******************************************************************************
#******************************************************************************
#******************************************************************************

# get the fluid speed that brings about a given specific pressure loss

def find_specific_pressure_loss_matching_flow(
        pipe: fic.Pipe,
        fluid: fic.Fluid,
        specific_pressure_loss: float,
        tolerance: float = 1e-6,
        dff_model = fic.DFF_COLEBROOK_WHITE,
        dff_model_is = fic.DFF_HAALAND,
        x0_fluid_speed = None,
        x0_reynolds_number = None,
        x0_friction_factor = None):
    """Returns the fluid speed in a pipe conducive to a given specific pressure loss."""

    # the pipe specific pressure drop appears to be a monotonic function of the
    # fluid speed, at least for each type of flow (laminar or turbulent).
    # if that is the case, a maximum pressure drop corresponds to a fluid speed,

    # assume the flow is laminar
    
    v_laminar = (specific_pressure_loss*
                 (pipe.d_int**2)/
                 (32*fluid.mass_density*fluid.kinematic_viscosity))
    
    # find the equivalent reynolds number
                 
    re_laminar = core.ReynoldsNumber(pipe.d_int, 
                                     v_laminar, 
                                     fluid.kinematic_viscosity)
    
    # friction factor
    
    f_laminar = 64/re_laminar
    
    # find out if the flow is laminar
    
    if (re_laminar <= fic.LAMINAR_UPPER_RE_LIMIT and
        re_laminar >= fic.LAMINAR_LOWER_RE_LIMIT):
    
        # the specific pressure loss is consistent with laminar flow
        
        # return values for laminar flow
    
        return v_laminar, re_laminar, f_laminar
    
    else:
        
        # turbulent or transition regime: it is necessary to solve a system

        #**********************************************************************
        #**********************************************************************
        
        # if the model is not recognised, use the colebrook-white one
        
        if dff_model not in fic.DFF_MODELS:
            
            dff_model = fic.DFF_COLEBROOK_WHITE
        
        # if the initial solution model is not recognised, use the haaland one
            
        if dff_model_is not in fic.DFF_EXPLICIT_MODELS:
            
            dff_model_is = fic.DFF_HAALAND
        
        # if there is no initial solution: get one
        
        if (x0_fluid_speed == None or 
            x0_friction_factor == None or
            x0_reynolds_number == None):
            
            # no initial solution has been provided
            
            # get an initial solution
            
            x0_reynolds_number = fic.DFF_RE_LIMITS[dff_model_is][1]
            
            x0_friction_factor = fic.DarcyFrictionFactor(x0_reynolds_number,
                                                         pipe,
                                                         model=dff_model_is)
            
            x0_fluid_speed = core.MeanFluidSpeedFromReynoldsNumber(
                pipe.d_int,
                x0_reynolds_number, 
                fluid.kinematic_viscosity)

        #**********************************************************************
        #**********************************************************************
        
        # system of equations:
        # 1) reynolds number
        # 2) friction factor equation
        # 3) specific pressure loss
        
        def system_equations(se):
            re, f, v = se
                
            # return (re-core.ReynoldsNumber(pipe.d_int, 
            #                                 v, 
            #                                 fluid.kinematic_viscosity),
            #         fic.DarcyFrictionFactorEquationTurbulentFlow(
            #                 re,
            #                 f,
            #                 pipe.e_eff,
            #                 pipe.d_int,
            #                 model=dff_model),
            #         specific_pressure_loss - 
            #             fic.SpecificPressureLossInPipe(
            #                 fluid.mass_density,
            #                 pipe.d_int,
            #                 v,
            #                 f))
            
            if re > 0 and f > 0 and v > 0:
                
                return (re-core.ReynoldsNumber(pipe.d_int, 
                                                v, 
                                                fluid.kinematic_viscosity),
                        fic.DarcyFrictionFactorEquationTurbulentFlow(
                                re,
                                f,
                                pipe.e_eff,
                                pipe.d_int,
                                model=dff_model),
                        specific_pressure_loss - 
                            fic.SpecificPressureLossInPipe(
                                fluid.mass_density,
                                pipe.d_int,
                                v,
                                f))
            
            else:
                
                out_re = 0
                out_f = 0
                out_v = 0
                
                if re <= 0:
                    
                    out_re = -math.inf
                    
                else:
                    
                    out_re = 0
                    
                if f <= 0:
                    
                    out_f = -math.inf
                    
                else:
                    
                    out_f = 0
                    
                if v <= 0:
                    
                    out_v = -math.inf
                    
                else:
                    
                    out_v = 0
                
                return (out_re, out_f, out_v)
        
        #**********************************************************************
        #**********************************************************************
        
        re, f, v =  fsolve(system_equations, 
                           (x0_reynolds_number,
                            x0_friction_factor,
                            x0_fluid_speed),
                           xtol=tolerance
                           )
        
        #**********************************************************************
        #**********************************************************************
        
        # return results
        
        return v, re, f
        
#******************************************************************************
#******************************************************************************
#******************************************************************************
#******************************************************************************
#******************************************************************************

@dataclass
class FlowRegimeInsidePipe:
    """A class for flow regimes inside pipes."""
    pipe: fic.Pipe
    fluid: fic.Fluid
    reynolds_number: float = None          
    fluid_speed: float = None
    mass_flow_rate: float = None          
    volumetric_flow_rate: float = None         
    
    def __post_init__(self):
    
        # at least one of the keyword arguments has to be non-default
        
        if (self.fluid_speed != None and
            self.mass_flow_rate == None and
            self.volumetric_flow_rate == None and 
            self.reynolds_number == None):
            
            self._use_fluid_speed()
        
            #******************************************************************
            #******************************************************************
            
        elif (self.fluid_speed == None and
              self.mass_flow_rate != None and
              self.volumetric_flow_rate == None and 
              self.reynolds_number == None):
            
            self._use_mass_flow_rate()
        
            #******************************************************************
            #******************************************************************
        
        elif (self.fluid_speed == None and
              self.mass_flow_rate == None and
              self.volumetric_flow_rate != None and 
              self.reynolds_number == None):
            
            self._use_volumetric_flow_rate()
        
            #******************************************************************
            #******************************************************************
        
        elif (self.fluid_speed == None and
              self.mass_flow_rate == None and
              self.volumetric_flow_rate == None and 
              self.reynolds_number != None):
            
            self._use_reynolds_number()
        
            #******************************************************************
            #******************************************************************
        
        else: 
            
            # the inputs were incorrectly specified
            
            raise NotImplementedError
    
            #******************************************************************
            #******************************************************************
        
        #**********************************************************************
        #**********************************************************************

    #**************************************************************************
    #**************************************************************************
    
    def _use_reynolds_number(self):
        
        # the reynolds number is known
         
        # fluid speed
        
        self.fluid_speed = core.MeanFluidSpeedFromReynoldsNumber(
            self.pipe.d_int, 
            self.reynolds_number, 
            self.fluid.kinematic_viscosity)
        
        # volumetric flow rate
        
        self.volumetric_flow_rate = core.VolumetricFlowRateFromMeanFluidSpeed(
            self.fluid_speed, 
            self.pipe.InternalCrossSectionalArea())
        
        # mass flow rate
        
        self.mass_flow_rate = core.MassFlowRateFromVolumetricFlowRate(
            self.volumetric_flow_rate, 
            self.fluid.mass_density)

    #**************************************************************************
    #**************************************************************************
    
    def _use_fluid_speed(self):
         
        # the fluid speed is known
        
        # volumetric flow rate
        
        self.volumetric_flow_rate = core.VolumetricFlowRateFromMeanFluidSpeed(
            self.fluid_speed, 
            self.pipe.InternalCrossSectionalArea())
        
        # mass flow rate
        
        self.mass_flow_rate = core.MassFlowRateFromVolumetricFlowRate(
            self.volumetric_flow_rate, 
            self.fluid.mass_density)
        
        # reynolds number
        
        self.reynolds_number = core.ReynoldsNumber(
            self.pipe.d_int, 
            self.fluid_speed, 
            self.fluid.kinematic_viscosity)

    #**************************************************************************
    #**************************************************************************
    
    def _use_volumetric_flow_rate(self):
  
        # the volumetric flow rate is known
          
        # fluid speed
        
        self.fluid_speed = core.MeanFluidSpeedViaVolumetricFlowRate(
            self.volumetric_flow_rate,
            self.pipe.InternalCrossSectionalArea())
        
        # mass flow rate
        
        self.mass_flow_rate = core.MassFlowRateFromVolumetricFlowRate(
            self.volumetric_flow_rate, 
            self.fluid.mass_density)
        
        # reynolds number
        
        self.reynolds_number = core.ReynoldsNumber(
            self.pipe.d_int, 
            self.fluid_speed, 
            self.fluid.kinematic_viscosity)
    
    #**************************************************************************
    #**************************************************************************
    
    def _use_mass_flow_rate(self):
        
        # the mass flow rate is known
    
        # volumetric flow rate
        
        self.volumetric_flow_rate = core.VolumetricFlowRateFromMassFlowRate(
            self.mass_flow_rate, 
            self.fluid.mass_density)
        
        # fluid speed
        
        self.fluid_speed = core.MeanFluidSpeedViaVolumetricFlowRate(
            self.volumetric_flow_rate,
            self.pipe.InternalCrossSectionalArea())
        
        # reynolds number
        
        self.reynolds_number = core.ReynoldsNumber(
            self.pipe.d_int, 
            self.fluid_speed, 
            self.fluid.kinematic_viscosity)

    #**************************************************************************
    #**************************************************************************

#******************************************************************************
#******************************************************************************
#******************************************************************************
#******************************************************************************
#******************************************************************************
#******************************************************************************